from django.db import models
from valar.models.core import VModel

class ForecastTask(VModel):
    """预测任务模型"""

    # Meta信息
    task_name = models.CharField(max_length=255, verbose_name="任务名称", blank=True, null=True)
    description = models.TextField(verbose_name="任务描述", blank=True, null=True)
    created_date = models.DateTimeField(auto_now_add=True, verbose_name="创建时间")
    updated_date = models.DateTimeField(auto_now=True, verbose_name="更新时间")

    # 预测参数
    forecast_horizon = models.IntegerField(verbose_name="预测步长", default=5)
    start_year = models.IntegerField(verbose_name="数据起始年份", null=True, blank=True)
    end_year = models.IntegerField(verbose_name="数据结束年份", null=True, blank=True)

    # 状态字段
    STATUS_CHOICES = [
        ('pending', '待处理'),
        ('running', '运行中'),
        ('completed', '已完成'),
        ('failed', '失败'),
    ]
    status = models.CharField(max_length=20, choices=STATUS_CHOICES, default='pending', verbose_name="任务状态")

    class Meta:
        verbose_name = '预测任务'
        verbose_name_plural = '预测任务'

    def __str__(self):
        return f"{self.task_name or '未命名任务'} - {self.created_date.strftime('%Y-%m-%d')}"


class TimeSeriesData(VModel):
    """时间序列数据模型"""

    task = models.ForeignKey(ForecastTask, on_delete=models.CASCADE, verbose_name="关联任务")

    # 时间和地区信息
    year = models.IntegerField(verbose_name="年份")
    province = models.CharField(max_length=100, verbose_name="省份")
    region = models.CharField(max_length=100, verbose_name="区域", blank=True, null=True)

    # 数据值
    value = models.FloatField(verbose_name="数值")

    # 数据类型标识
    DATA_TYPE_CHOICES = [
        ('original', '原始数据'),
        ('proportion', '比例数据'),
        ('ilr_transformed', 'ILR变换数据'),
    ]
    data_type = models.CharField(max_length=20, choices=DATA_TYPE_CHOICES, default='original', verbose_name="数据类型")

    class Meta:
        verbose_name = '时间序列数据'
        verbose_name_plural = '时间序列数据'
        unique_together = ['task', 'year', 'province', 'data_type']

    def __str__(self):
        return f"{self.province} - {self.year}: {self.value}"


class ForecastResult(VModel):
    """预测结果模型"""

    task = models.ForeignKey(ForecastTask, on_delete=models.CASCADE, verbose_name="关联任务")

    # 预测框架类型
    FRAMEWORK_CHOICES = [
        ('bottom_up', '自底向上预测'),
        ('mint', 'MinT协调预测'),
        ('ilr_var', 'ILR-VAR预测'),
    ]
    framework = models.CharField(max_length=20, choices=FRAMEWORK_CHOICES, verbose_name="预测框架")

    # 预测时间和地区
    forecast_year = models.IntegerField(verbose_name="预测年份")
    province = models.CharField(max_length=100, verbose_name="省份")
    region = models.CharField(max_length=100, verbose_name="区域", blank=True, null=True)

    # 预测值和置信区间
    forecast_value = models.FloatField(verbose_name="预测值")
    lower_bound = models.FloatField(verbose_name="置信区间下界", null=True, blank=True)
    upper_bound = models.FloatField(verbose_name="置信区间上界", null=True, blank=True)

    # 预测结果生成时间
    created_date = models.DateTimeField(auto_now_add=True, verbose_name="生成时间")

    class Meta:
        verbose_name = '预测结果'
        verbose_name_plural = '预测结果'
        unique_together = ['task', 'framework', 'forecast_year', 'province']

    def __str__(self):
        return f"{self.framework} - {self.province} - {self.forecast_year}: {self.forecast_value}"


class HierarchicalStructure(VModel):
    """层级结构配置模型"""

    name = models.CharField(max_length=100, unique=True, verbose_name="结构名称")
    description = models.TextField(verbose_name="结构描述", blank=True, null=True)

    # 层级配置JSON
    hierarchy_config = models.JSONField(verbose_name="层级配置", default=dict)

    # 默认配置标识
    is_default = models.BooleanField(verbose_name="是否为默认配置", default=False)

    class Meta:
        verbose_name = '层级结构'
        verbose_name_plural = '层级结构'

    def __str__(self):
        return self.name

    def get_regions(self):
        """获取区域配置"""
        return self.hierarchy_config.get('regions', {})

    def get_provinces_by_region(self, region):
        """获取指定区域的省份列表"""
        return self.hierarchy_config.get('regions', {}).get(region, [])


